# utils function
def is_in_network(ip, network):
    # 简单示例，假设网络格式为 'xxx.xxx.xxx.0/24'
    network_ip, prefix = network.split('/')
    network_parts = list(map(int, network_ip.split('.')))
    ip_parts = list(map(int, ip.split('.')))
    mask = (0xFFFFFFFF << (32 - int(prefix))) & 0xFFFFFFFF
    network_int = (network_parts[0] << 24) | (network_parts[1] << 16) | (network_parts[2] << 8) | network_parts[3]
    ip_int = (ip_parts[0] << 24) | (ip_parts[1] << 16) | (ip_parts[2] << 8) | ip_parts[3]
    return (ip_int & mask) == (network_int & mask)

def get_network_from_ip(ip_mask):
    ip, prefix = ip_mask.split('/')
    ip_parts = list(map(int, ip.split('.')))
    mask = (0xFFFFFFFF << (32 - int(prefix))) & 0xFFFFFFFF
    network_int = (ip_parts[0] << 24) | (ip_parts[1] << 16) | (ip_parts[2] << 8) | ip_parts[3]
    network_ip = (network_int & mask)
    network_parts = [
        (network_ip >> 24) & 0xFF,
        (network_ip >> 16) & 0xFF,
        (network_ip >> 8) & 0xFF,
        network_ip & 0xFF
    ]
    return f"{network_parts[0]}.{network_parts[1]}.{network_parts[2]}.{network_parts[3]}/{prefix}"

# 定义 Packet 类，包含二层、三层和自定义数据
class Packet:
    def __init__(self, src_mac, dst_mac, src_ip, dst_ip, ip_type='IP', data=''):
        # L2 info
        self.src_mac = src_mac
        self.dst_mac = dst_mac
        # L3 info
        self.src_ip = src_ip
        self.dst_ip = dst_ip
        # PDU type
        self.ip_type = ip_type
        self.data = data


# 定义 Port 类
class Port:
    def __init__(self, device, index, mac_address=None, ip_mask=None):
        self.device = device
        self.index = index
        self.mac_address = mac_address
        self.ip_address = None
        if ip_mask:
            self.ip_address, self.prefix = ip_mask.split('/')
            self.network = get_network_from_ip(ip_mask)
        self.link = None
        # 等待发送的数据包列表{arp_dst_ip: [pkt1, pkt2]}（比如未获取目的mac的数据包，需要先缓存，等arp响应到达之后再发送）
        self.waiting_packets = {}

    def connect_link(self, link):
        self.link = link


# 定义一个抽象的网络设备类
class NetworkDevice:
    def __init__(self, name, mac_address=None, ip_mask=None):
        self.name = name
        # 接口列表
        self.ports = {}
        # 路由表，格式：{网段: (下一跳 IP, 类型)}
        self.routing_table = {}
        if mac_address and ip_mask:
            self.add_port('eth1', mac_address, ip_mask)

    def add_port(self, index, mac_address=None, ip_mask=None):
        port = Port(self, index, mac_address, ip_mask)
        self.ports[index] = port
        if port.ip_address:
            # 添加直连路由
            self.add_route(port.network, port.ip_address, '直连')
        return port

    def add_route(self, network, next_hop, route_type='静态'):
        print(f'[{self.name}] add route {network} -> {next_hop} ({route_type})')
        self.routing_table[network] = (next_hop, route_type)

    def send_packet(self, packet, out_port_index=None):
        print('no implement send_packet')

    def get_route_for_ip(self, ip):
        for network, (next_hop, route_type) in self.routing_table.items():
            if is_in_network(ip, network):
                return next_hop, route_type
        return None, None


# 定义交换机类，继承自 NetworkDevice
class Switch(NetworkDevice):
    def __init__(self, name):
        super().__init__(name)
        self.mac_table = {}

    def receive_packet(self, packet, in_port):
        # L2 switch logic
        # 1. 学习源 MAC 地址
        self.mac_table[packet.src_mac] = in_port.index
        print(f"交换机 {self.name} 学习到 MAC 地址 {packet.src_mac} 位于端口 {in_port.index}")

        # 2. 如果目的 MAC 地址在 MAC 地址表中，将数据包转发到对应的端口
        if packet.dst_mac in self.mac_table:
            target_port_index = self.mac_table[packet.dst_mac]
            target_port = self.ports.get(target_port_index)
            print(f"交换机 {self.name} 将数据包从端口 {in_port.index} 转发到端口 {target_port_index}")
            target_port.link.transmit(packet, target_port)
        # 3. 如果目的 MAC 地址不在 MAC 地址表中，将数据包广播到所有端口
        else:
            print(f"交换机 {self.name} 未找到 MAC 地址 {packet.dst_mac}，进行广播")
            for port_index, port in self.ports.items():
                # 跳过输入端口(避免环路)，和未连接的端口
                if port != in_port and port.link:
                    print(f'{self.name} send packet to port {port_index}')
                    port.link.transmit(packet, port)

    def send_packet(self, packet, out_port_index):
        print('switch itself should not send packet')

# 定义主机类，继承自 NetworkDevice
class Host(NetworkDevice):
    def __init__(self, name, mac_address, ip_mask):
        super().__init__(name, mac_address, ip_mask)
        # 默认路由下一跳
        self.default_gateway = None
        # ARP 表，格式：{IP 地址: MAC 地址}
        self.arp_table = {}
        self.mac_address = mac_address
        self.ip_address = ip_mask.split('/')[0] if ip_mask else None

    def set_default_gateway(self, gateway_ip):
        self.default_gateway = gateway_ip

    def add_route(self, network, next_hop, route_type='静态'):
        super().add_route(network, next_hop, route_type)

    def send_packet(self, packet):
        default_port = self.ports['eth1']
        # 1. 查找路由表
        next_hop, route_type = self.get_route_for_ip(packet.dst_ip)
        print(f'lookup route for {packet.dst_ip}: route_type: {route_type}')
        if route_type == '直连':
            print(f'主机 {self.ip_address} 目的IP {packet.dst_ip} 在直连网段，查找目的主机ARP')
            arp_dst_ip = packet.dst_ip
        elif route_type == '静态':
            print(f'主机 {self.ip_address} 开始查找路由下一跳ARP')
            arp_dst_ip = next_hop
        else:
            print(f"通过默认网关({self.default_gateway}) 发送，查找默认网关ARP")
            arp_dst_ip = self.default_gateway
        # 查看ARP表
        dst_mac = self.arp_table.get(arp_dst_ip)
        if not dst_mac:
            print(f"主机 {self.ip_address} ARP 表中没有 {arp_dst_ip} 的 MAC 地址，发送 ARP 请求")
            # 缓存IP报文
            if arp_dst_ip not in default_port.waiting_packets:
                default_port.waiting_packets[arp_dst_ip]= []
            default_port.waiting_packets[arp_dst_ip].append(packet)
            arp_packet = Packet(
                src_mac=self.mac_address,
                dst_mac="FF:FF:FF:FF:FF:FF",
                src_ip=self.ip_address,
                dst_ip=arp_dst_ip,
                ip_type="ARP Request"
            )
            if default_port.link:
                default_port.link.transmit(arp_packet, default_port)
        else:
            packet.dst_mac = dst_mac
            print(f"主机 {self.ip_address} 查找ARP表，将数据包发往（MAC: {dst_mac}）")
            default_port.link.transmit(packet, default_port)

    def receive_packet(self, packet, in_port):
        print(f"[{self.name}]:{in_port.index} 收到 {packet.ip_type}[]")
        if packet.dst_ip == self.ip_address:
            if 'ARP' in packet.ip_type:
                if packet.ip_type == "ARP Request":
                    print(f"主机 {self.ip_address} 收到 ARP 请求，发送 ARP 响应")
                    arp_response_packet = Packet(
                        src_mac=self.mac_address,
                        dst_mac=packet.src_mac,
                        src_ip=self.ip_address,
                        dst_ip=packet.src_ip,
                        ip_type="ARP Response"
                    )
                    in_port.link.transmit(arp_response_packet, in_port)
                elif packet.ip_type == "ARP Response":
                    print(f"主机 {self.ip_address} 收到 ARP 响应，更新 ARP 表")
                    self.arp_table[packet.src_ip] = packet.src_mac
                    # 如果有等待此arp的IP报文，则发送
                    for waiting_packet in in_port.waiting_packets.get(packet.src_ip, []):
                        print(f"主机 {self.ip_address} 发送等待的 {waiting_packet.ip_type} 数据包")
                        waiting_packet.dst_mac = packet.src_mac
                        in_port.link.transmit(waiting_packet, in_port)
            else:
                print(f"主机 {self.ip_address} 接收到来自 {packet.src_ip} 的数据包，自定义数据: {packet.data}")
                # 处理数据包
        else:
            print(f"主机 {self.ip_address} 接收到非目的数据包，丢弃")


# 定义路由器类，继承自 NetworkDevice
class Router(NetworkDevice):
    def __init__(self, name, mac_address, ip_address):
        super().__init__(name, mac_address, ip_address)
        # 路由器的路由表，格式：{网段: (下一跳 IP, 类型)}
        self.routing_table = {}
        # ARP 表，格式：{IP 地址: MAC 地址}
        self.arp_table = {}

    def add_route(self, network, next_hop, route_type='静态'):
        super().add_route(network, next_hop, route_type)

    def send_packet(self, packet, in_port):
        # 1. 查找路由表
        next_hop, route_type = self.get_route_for_ip(packet.dst_ip)
        out_port = None
        print(f'[{self.name}] 1. lookup route for {packet.dst_ip}: next_hop: {next_hop}')
        # 2. 根据路由类型找out_port
        if not next_hop:
            print(f"路由器 {self.name} 没有找到合适的路由，丢弃数据包")
            # todo: reply icmp message(network unreachable)
        if route_type == '直连':
            print(f"路由器 {self.name} 发现直连网段，直接发送数据包")
            out_port = [port for port in self.ports.values() if port.ip_address == next_hop][0]
            arp_dst_ip = packet.dst_ip
        elif route_type == '静态':
            print(f"路由器 {self.name} 发现静态路由，根据下一跳IP({next_hop})找出接口")
            #
            out_port = [port for port in self.ports.values() if is_in_network(next_hop, port.network)]
            if not out_port:
                print(f"路由器 {self.name} 下一跳network不可达，丢弃数据包")
                return
            arp_dst_ip = next_hop
        # 3. 查找下一跳IP的ARP表
        packet.src_mac = out_port.mac_address
        dst_mac = self.arp_table.get(arp_dst_ip)
        if not dst_mac:
            # cache pkt
            if arp_dst_ip not in out_port.waiting_packets:
                out_port.waiting_packets[arp_dst_ip] = []
            out_port.waiting_packets[arp_dst_ip].append(packet)
            print(f"路由器 {self.name} ARP 表中没有 {arp_dst_ip} 的 MAC 地址，发送 ARP 请求")
            arp_packet = Packet(
                src_mac=out_port.mac_address,
                dst_mac="FF:FF:FF:FF:FF:FF",
                src_ip=out_port.ip_address,
                dst_ip=arp_dst_ip,
                ip_type="ARP Request"
            )
            out_port.link.transmit(arp_packet, out_port)
        else:
            packet.dst_mac = dst_mac
            print(f"路由器 {self.name} 查找路由表，将数据包发往 {arp_dst_ip}（MAC: {dst_mac}）")
            out_port.link.transmit(packet, out_port)

    def receive_packet(self, packet, in_port):
        print(f"[{self.name}]:{in_port.index} recv {packet.ip_type}：[{packet.src_ip}->{packet.dst_ip}]")
        # 1. 处理local packet
        if packet.dst_ip == in_port.ip_address:
            if packet.ip_type == "ARP Request":
                print(f"[{self.name}] 收到 ARP 请求，发送 ARP 响应")
                arp_response_packet = Packet(
                    src_mac=in_port.mac_address,
                    dst_mac=packet.src_mac,
                    src_ip=packet.dst_ip,
                    dst_ip=packet.src_ip,
                    ip_type="ARP Response"
                )
                in_port.link.transmit(arp_response_packet, in_port)
            elif packet.ip_type == "ARP Response":
                print(f"[{self.name}] 收到 ARP 响应，更新 ARP 表")
                self.arp_table[packet.src_ip] = packet.src_mac
                # 如果有等待此arp的IP报文，则发送
                for waiting_packet in in_port.waiting_packets.get(packet.src_ip, []):
                    print(f"[{self.name}]:{in_port.index} 发送等待的 {waiting_packet.ip_type} 数据包")
                    waiting_packet.dst_mac = packet.src_mac
                    in_port.link.transmit(waiting_packet, in_port)
            else:
                print(f"[{self.name}] 接收到来自 {packet.src_ip} 的数据包，自定义数据: {packet.data}")
                # todo: icmp reply
        else:
            if packet.ip_type == "ARP Request":
                print(f"[{self.name}] 收到非本设备 ARP 请求，drop")
                return
            # 非本机报文，通过路由转发
            self.send_packet(packet, in_port)


# 定义链路类
class Link:
    def __init__(self, port1, port2):
        self.port1 = port1
        self.port2 = port2
        port1.connect_link(self)
        port2.connect_link(self)

    def transmit(self, packet, in_port):
        if in_port not in [self.port1, self.port2]:
            print('链路错误: 端口不在链路内')
        if in_port == self.port1:
            target_port = self.port2
        else:
            target_port = self.port1
        target_port.device.receive_packet(packet, target_port)


if __name__ == '__main__':
    # 创建设备实例
    host1 = Host('host1', "00:11:22:33:44:11", "192.168.1.100/24")
    host2 = Host('host2', "00:11:22:33:44:22", "192.168.1.200/24")
    host3 = Host('host3', "00:11:22:33:44:33", "192.168.2.100/24")
    switch = Switch('sw1')
    router = Router('router1', "AA:22:33:44:55:11", "192.168.1.254/24")

    # 添加端口和接口
    switch_port1 = switch.add_port(1)
    switch_port2 = switch.add_port(2)
    switch_port3 = switch.add_port(3)
    router_port2 = router.add_port('eth2', "AA:22:33:44:55:22", "192.168.2.254/24")

    # 创建链路实例
    link1 = Link(switch_port1, host1.ports['eth1'])
    link2 = Link(switch_port2, host2.ports['eth1'])
    link3 = Link(switch_port3, router.ports['eth1'])
    link4 = Link(router_port2, host3.ports['eth1'])

    # 配置主机 1 的路由表和默认网关
    host1.set_default_gateway("192.168.1.254")
    host2.set_default_gateway("192.168.2.254")

    # 模拟主机 1 向主机 2 发送数据包
    # --创建一个 Packet 实例
    packet = Packet(
        src_mac=host1.mac_address,
        dst_mac="unknown",
        src_ip=host1.ports['eth1'].ip_address,
        dst_ip=host2.ports['eth1'].ip_address,
        ip_type="Hello, Host 2!"
    )
    print("同网段测试：host1 ping host2...")
    #host1.send_packet(packet)
    print("不同网段测试：host1 ping host3...")
    packet = Packet(
        src_mac=host1.mac_address,
        dst_mac="unknown",
        src_ip=host1.ports['eth1'].ip_address,
        dst_ip=host3.ports['eth1'].ip_address,
        ip_type='IP', 
        data="Hello, Host 3!"
    )
    host1.send_packet(packet)